#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created by PyCharm.

@Date    : Wed Feb 24 2021 
@Time    : 11:04:37
@File    : dataset.py
@Author  : alpha
"""


import random

import cv2
import torch
import numpy as np
import matplotlib.pyplot as plt
import albumentations as AT
import torchvision.transforms as TT

from pathlib import Path
from PIL import Image
from torch.utils.data import Dataset

from src.log import logger


class CelebASpoofCroppedDataset(Dataset):

    def __init__(self, root, mode='train', transform=None):
        super(CelebASpoofCroppedDataset, self).__init__()
        assert mode in ['train', 'test']
        self.data_path = Path(root) / mode

        logger.info('Loading the cropped CelebASpoof dataset {} indices...'.format(mode))
        self.datalist_spoof = sorted(self.data_path.glob('spoof/**/*.jpg'))
        self.datalist_live = sorted(self.data_path.glob('live/**/*.jpg'))
        self.datalist = self._get_datalist_with_label()
        logger.info('DONE!')

        self.transform = transform
        self.to_tensor = TT.Compose([
            TT.ToTensor(),
            TT.Normalize(mean=(0.5, 0.5, 0.5), std=(1., 1., 1.))
        ])

    def _get_datalist_with_label(self):
        list_spoof_with_label = list(zip(self.datalist_spoof, np.ones(len(self.datalist_spoof), dtype=np.int)))
        list_live_with_label = list(zip(self.datalist_live, np.zeros(len(self.datalist_live), dtype=np.int)))
        return list_live_with_label + list_live_with_label + list_spoof_with_label

    def __getitem__(self, item):
        # anchor, pos, neg, label
        anchor_im_path, label = self.datalist[item]
        if label == 0:
            pos_im_path = random.choice(self.datalist_live)
            neg_im_path = random.choice(self.datalist_spoof)
        else:
            pos_im_path = random.choice(self.datalist_spoof)
            neg_im_path = random.choice(self.datalist_live)

        anchor_im = np.array(Image.open(anchor_im_path))
        pos_im = np.array(Image.open(pos_im_path))
        neg_im = np.array(Image.open(neg_im_path))

        if self.transform is not None:
            anchor_im = self.transform(image=anchor_im)['image']
            pos_im = self.transform(image=pos_im)['image']
            neg_im = self.transform(image=neg_im)['image']

        anchor_im = self.to_tensor(anchor_im)
        pos_im = self.to_tensor(pos_im)
        neg_im = self.to_tensor(neg_im)
        label = torch.tensor(label, dtype=torch.long)

        return anchor_im, pos_im, neg_im, label

    def __len__(self):
        return len(self.datalist)


class FaceSpoofEvalDataset(Dataset):

    def __init__(self, root, exts=['.jpg', '.png']):
        super(FaceSpoofEvalDataset, self).__init__()
        self.datalist = sorted(Path(root).glob('**/*[{}]'.format('|'.join(exts))))
        self.to_tensor = TT.Compose([
            TT.Resize((256, 256)),
            TT.ToTensor(),
            TT.Normalize(mean=(0.5, 0.5, 0.5), std=(1., 1., 1.))
        ])


    def __getitem__(self, item):
        bgr = cv2.imread(str(self.datalist[item]))
        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
        im = Image.fromarray(rgb)
        # im = Image.open(self.datalist[item], mode='RGB')
        return self.to_tensor(im), self.datalist[item]

    def __len__(self):
        return len(self.datalist)


if __name__ == '__main__':

    transform = AT.Compose([
        AT.HorizontalFlip(p=0.5),
        AT.MotionBlur(p=0.2),
        AT.Rotate(limit=20, p=0.2),
        AT.RandomBrightnessContrast(p=0.2),
        AT.RandomResizedCrop(256, 256, scale=(0.8, 1.0), ratio=(1, 1), p=0.99)
    ])
    dataset = CelebASpoofCroppedDataset(
        root='/data/dataset/CelebA_Spoof_cropped', mode='test', transform=transform
    )
    for idx in range(5):
        anchor_im, pos_im, neg_im, label = dataset[idx]
        aim = ((anchor_im + 0.5).permute((1, 2, 0)).data.numpy() * 255).astype(np.uint8)
        pim = ((pos_im + 0.5).permute((1, 2, 0)).data.numpy() * 255).astype(np.uint8)
        nim = ((neg_im + 0.5).permute((1, 2, 0)).data.numpy() * 255).astype(np.uint8)

        plt.subplots(nrows=1, ncols=3)
        plt.subplot(131)
        plt.imshow(aim)
        plt.title('label: {}'.format(label))
        plt.subplot(132)
        plt.imshow(pim)
        plt.subplot(133)
        plt.imshow(nim)
        plt.tight_layout()
        plt.show()